﻿from towhee import ops, pipe

class SentenceModel:
    """
    SentenceModel
    """

    def __init__(self):
        self.sentence_embedding_pipe = (
            pipe.input('sentence')
                .map('sentence', 'embedding', ops.text_embedding.transformers(model_name='text2vec-base-chinese'))
                .map('embedding', 'embedding', lambda x: x[0])
                .map('embedding', 'embedding', ops.towhee.np_normalize())
                .output('embedding')
        )

    def sentence_encode(self, data_list):
        res_list = []
        for data in data_list:
            res = self.sentence_embedding_pipe(data)
            res_list.append(res.get()[0])

        return res_list


if __name__ == '__main__':
    MODEL = SentenceModel()
    d1=["你好"]
    #r1=m.encode(d1)
    #rr=list(map(lambda x:ops.towhee.np_normalize(x),r1))
    #print(rr)
    # Warm up the model to build image
    t1=MODEL.sentence_encode(d1)
    print(t1[0])
